import datetime
import pandas as pd
import finplot as fplt
from gofishing.base import TradeType

fplt.display_timezone = datetime.timezone.utc
fplt.candle_bull_color, fplt.candle_bear_color = '#ef5350', '#26a69a'
fplt.volume_bull_color, fplt.volume_bear_color = '#f7a9a7', '#92d2cc'
fplt.volume_bull_body_color = '#f7a9a7'


def plotKLine(pv: pd.DataFrame) -> None:
    KLine = pv
    KLine.reset_index(inplace=True)
    ax = fplt.create_plot('K-Line plot')
    fplt.candlestick_ochl(KLine[['time', 'open', 'close', 'high', 'low']],
                          ax=ax)
    fplt.volume_ocv(KLine[['time', 'open', 'close', 'volume']],
                    ax=ax.overlay())
    fplt.show()


def plotTrade(pv: pd.DataFrame, tradeRecord: pd.DataFrame) -> None:

    timeIndex = pv.index.copy()

    def _alignForPlot(df: pd.DataFrame) -> pd.DataFrame:
        return df.set_index('time').reindex(timeIndex).reset_index()

    KLine = pv.copy()
    KLine.reset_index(inplace=True)

    tradeRecord = tradeRecord.copy()
    tradeRecord['time'] = pd.Timestamp(0)
    tradeTime = tradeRecord['tradeTime']

    # 将 tradeTime 与 timeIndex 对齐
    for indice in range(len(tradeRecord)):
        for time in timeIndex:
            if tradeTime[indice] <= time:
                tradeRecord.at[indice, 'time'] = time
                break

    ax = fplt.create_plot('K-Line plot', rows=2)
    fplt.candlestick_ochl(KLine[['time', 'open', 'close', 'high', 'low']],
                          ax=ax[0])
    fplt.volume_ocv(KLine[['time', 'open', 'close', 'volume']],
                    ax=ax[0].overlay())

    colorTrade = {TradeType.LONGOPEN: 'b', TradeType.LONGCLOSE: 'b',
                  TradeType.SHORTOPEN: 'b', TradeType.SHORTCLOSE: 'b'}
    styleTrade = {TradeType.LONGOPEN: '^', TradeType.LONGCLOSE: 'x',
                  TradeType.SHORTOPEN: 'v', TradeType.SHORTCLOSE: 'x'}
    for tradeType in TradeType:
        tradeMark = tradeRecord[tradeRecord['tradeType'] == tradeType]
        if not tradeMark.empty:
            tradeMark = _alignForPlot(tradeMark[['time', 'tradePrice']])
            fplt.plot(tradeMark['time'], tradeMark['tradePrice'], ax=ax[0],
                      color=colorTrade[tradeType], style=styleTrade[tradeType],
                      legend=tradeType.value)

    fplt.show()
